{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90ad74ee",
   "metadata": {},
   "source": [
    "# Module & Pytree"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0691d0",
   "metadata": {},
   "source": [
    "Flax NNX's Modules are by default registered as JAX Pytrees, this allows using them throughout most of JAX APIs but in particular JAX transforms and the `jax.tree.*` functions. Thanks to the pytree protocol a simple NNX program might look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9b2b929d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (5, 3)\n"
     ]
    }
   ],
   "source": [
    "from flax import nnx\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "class Linear(nnx.Module):\n",
    "  def __init__(self, din, dout, rngs: nnx.Rngs):\n",
    "    self.din, self.dout = din, dout\n",
    "    self.kernel = nnx.Param(rngs.normal((din, dout)))\n",
    "\n",
    "rngs = nnx.Rngs(0)\n",
    "weights = Linear(2, 3, rngs=rngs)\n",
    "\n",
    "@jax.jit\n",
    "def forward(weights, x):\n",
    "  return x @ weights.kernel\n",
    "\n",
    "y = forward(weights, x=rngs.uniform((5, 2)))\n",
    "print(f\"{y.shape = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcfbbdb2",
   "metadata": {},
   "source": [
    "Here `weights`, of type `Linear`, was able to be passed directly to the `jit`-ed function `forward`. Throughout the rest of this guide we will try to answer the questions:\n",
    "1. What are pytrees? \n",
    "2. How does NNX implement pytrees?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f610bb3",
   "metadata": {},
   "source": [
    "## Pytrees 101\n",
    "Most modern ML models have too many Arrays for users to pass around individually, to deal with this JAX developed a way to track Array data in nested structures that still allowed caching for compilation: Pytrees. JAX pytrees are tree structures made of python objects that can be recursively traversed in order to collect an ordered list of leaves and a definition of the tree structure, this is done via the `jax.tree.flatten` function. Most common pytrees are native python containers like `list`, `dict`, and `tuple`, but interestingly it also include `None`. The example bellow shows how to collect all the integer leaves from a nested structure using `flatten`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c3529274",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "leaves = [1, 2, 3, 4]\n",
      "treedef = PyTreeDef([{'a': *}, {'b': *, 'c': (*, *), 'd': None}])\n"
     ]
    }
   ],
   "source": [
    "pytree = [\n",
    "  {'a': 1},\n",
    "  {\n",
    "    'b': 2,\n",
    "    'c': (3, 4),\n",
    "    'd': None,\n",
    "  }\n",
    "]\n",
    "\n",
    "leaves, treedef = jax.tree.flatten(pytree)\n",
    "print(f\"leaves = {leaves}\")\n",
    "print(f\"treedef = {treedef}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c037b5e",
   "metadata": {},
   "source": [
    "Note that `None` is not a leaf because its defined as a pytree with no children. The main purpose of being able to flatten, apart from collecting the leaves, is being able reconstruct the pytree structure from the tree definition from any sequence of leaves of the same length via the `jax.tree.unflatten` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8237524",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "old pytree = [{'a': 1}, {'b': 2, 'c': (3, 4), 'd': None}]\n",
      "new pytree = [{'a': 10}, {'b': 20, 'c': (30, 40), 'd': None}]\n"
     ]
    }
   ],
   "source": [
    "new_leaves = [x * 10 for x in leaves]\n",
    "new_pytree = jax.tree.unflatten(treedef, new_leaves)\n",
    "\n",
    "print(f\"old pytree = {pytree}\")\n",
    "print(f\"new pytree = {new_pytree}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "117c8b2d",
   "metadata": {},
   "source": [
    "### Custom Pytrees\n",
    "JAX allows us to register custom pytree node type by using the `jax.tree_util.register_pytree_node` utility. For any type we are able to define a flatten that decomposes the object into a a sequence of nodes / children and a static (hashable) structure, and a unflatten function which takes the sequence of nodes and the static structure and creates a new instance. In the following example we create a simple type `Foo` with the attributes `a`, `b`, and `c`, and define `a` and `b` as nodes and `c` as static."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0c46905c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "leaves = [1, 2]\n",
      "treedef = PyTreeDef(CustomNode(Foo[('hi',)], [*, *]))\n"
     ]
    }
   ],
   "source": [
    "class Foo:\n",
    "  def __init__(self):\n",
    "    self.a = 1\n",
    "    self.b = 2\n",
    "    self.c = \"hi\"\n",
    "\n",
    "def flatten_foo(foo: Foo):\n",
    "  nodes = [foo.a, foo.b]  # sequence of nodes\n",
    "  static = (foo.c,) # hashable & equatable structure\n",
    "  return nodes, static\n",
    "\n",
    "def unflatten_foo(static, nodes):\n",
    "  foo = object.__new__(Foo)  # create uninitialized instance\n",
    "  foo.a = nodes[0]\n",
    "  foo.b = nodes[1]\n",
    "  foo.c = static[0]\n",
    "  return foo\n",
    "\n",
    "jax.tree_util.register_pytree_node(Foo, flatten_foo, unflatten_foo)\n",
    "\n",
    "foo = Foo()\n",
    "leaves, treedef = jax.tree.flatten(foo)\n",
    "print(f\"leaves = {leaves}\")\n",
    "print(f\"treedef = {treedef}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15647b13",
   "metadata": {},
   "source": [
    "Notice that `'hi'` does not appear in the leaves because `c` is defined as static, but you can see it as part of the `PyTreeDef` structure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d603fa09",
   "metadata": {},
   "source": [
    "## nnx.Pytree\n",
    "In general it would be cumbersome for users to manually register the pytree definition for every type they create. To automate this process NNX provides the `nnx.Pytree` base type that offers a simple API: users annotate attributes using either `nnx.static` or `nnx.data`, and Pytree will register some flatten and unflatten functions that will take the annotations into account. The `nnx.data` and `nnx.static` annotations must only be assigned to `Pytree` attributes directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "95016a94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pytree structure:\n",
      " - pytree.layers[0].b = Array([0.], dtype=float32)\n",
      " - pytree.layers[0].w = Array([[1.]], dtype=float32)\n",
      " - pytree.layers[1].b = Array([0.], dtype=float32)\n",
      " - pytree.layers[1].w = Array([[1.]], dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "class Linear(nnx.Pytree):\n",
    "  def __init__(self, din: int, dout: int):\n",
    "    self.din = nnx.static(din)\n",
    "    self.dout = nnx.static(dout)\n",
    "    self.w = nnx.data(jnp.ones((din, dout)))\n",
    "    self.b = nnx.data(jnp.zeros((dout,)))\n",
    "\n",
    "class MLP(nnx.Pytree):\n",
    "  def __init__(self, num_layers, dim):\n",
    "    self.num_layers = nnx.static(num_layers)\n",
    "    self.layers = nnx.data([\n",
    "      Linear(dim, dim) for _ in range(num_layers)\n",
    "    ])\n",
    "\n",
    "pytree = MLP(num_layers=2, dim=1)\n",
    "\n",
    "def pytree_structure(pytree, title='pytree structure'):\n",
    "  print(f\"{title}:\")\n",
    "  path_leaves, treedef = jax.tree.flatten_with_path(pytree)\n",
    "  for path, value in path_leaves:\n",
    "    print(f\" - pytree{jax.tree_util.keystr(path)} = {value!r}\")\n",
    "\n",
    "pytree_structure(pytree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a2db214",
   "metadata": {},
   "source": [
    "As you can see above, only the `data` paths appear in the leaves. However, its very verbose to have to define `static` and `data` for each attribute, so `Pytree` has sensible defaults. You can remove most of them and it will just work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a8665146",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pytree structure:\n",
      " - pytree.layers[0].b = Array([0.], dtype=float32)\n",
      " - pytree.layers[0].w = Array([[1.]], dtype=float32)\n",
      " - pytree.layers[1].b = Array([0.], dtype=float32)\n",
      " - pytree.layers[1].w = Array([[1.]], dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "class Linear(nnx.Pytree):\n",
    "  def __init__(self, din: int, dout: int):\n",
    "    self.din = din # static\n",
    "    self.dout = dout # static\n",
    "    self.w = jnp.ones((din, dout)) # data\n",
    "    self.b = jnp.zeros((dout,)) # data\n",
    "\n",
    "class MLP(nnx.Pytree):\n",
    "  def __init__(self, num_layers, dim):\n",
    "    self.num_layers = num_layers # static\n",
    "    self.layers = nnx.List([ # data\n",
    "      Linear(dim, dim) for _ in range(num_layers)\n",
    "    ])\n",
    "\n",
    "pytree = MLP(num_layers=2, dim=1)\n",
    "pytree_structure(pytree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a249c5f0",
   "metadata": {},
   "source": [
    "The only change we had to do here is use `nnx.List` to signal that `layers` contains `data`, the status of the rest of the attributes can be correctly inferred. The rules that determine if a value is data or not are the following:\n",
    "\n",
    "* `Array`s, `Variable`s, `ArrayRef`s, and `nnx.Pytree`s are data.\n",
    "* Types registered using `nnx.register_data_type` are data.\n",
    "* All other types are static.\n",
    "\n",
    "To check if a value is data use the `nnx.is_data` function which will return its status:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "27682936",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "# ------ DATA ------------\n",
      "nnx.is_data( jnp.array(0) ) = True                   # Arrays are data\n",
      "nnx.is_data( nnx.Param(1) ) = True                   # Variables are data\n",
      "nnx.is_data( nnx.Rngs(2) ) = True                    # nnx.Pytrees are data\n",
      "\n",
      "# ------ STATIC ------------\n",
      "nnx.is_data( 'hello' ) = False                       # strings, arbitrary objects\n",
      "nnx.is_data( 42 ) = False                            # int, float, bool, complex, etc.\n",
      "nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) = False    # list, dict, tuple, regular pytrees\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"\"\"\n",
    "# ------ DATA ------------\n",
    "{nnx.is_data( jnp.array(0) ) = }                   # Arrays are data\n",
    "{nnx.is_data( nnx.Param(1) ) = }                   # Variables are data\n",
    "{nnx.is_data( nnx.Rngs(2) ) = }                    # nnx.Pytrees are data\n",
    "\n",
    "# ------ STATIC ------------\n",
    "{nnx.is_data( 'hello' ) = }                       # strings, arbitrary objects\n",
    "{nnx.is_data( 42 ) = }                            # int, float, bool, complex, etc.\n",
    "{nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) = }    # list, dict, tuple, regular pytrees\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ab06e79",
   "metadata": {},
   "source": [
    "### When to use explicit annotations?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dead73d",
   "metadata": {},
   "source": [
    "There are cases were you do want to explicitely annotate the attributes to avoid ambiguity or protect yourself against possible edge cases. These include constraining input arguments which might have unexpected types, forcing attributes as data when their type is not treated as data by default, or using `nnx.static` as a way to assert the attribute should not contain data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9e064461",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pytree structure:\n",
      " - pytree.bias.value = Array(0., dtype=float32, weak_type=True)\n",
      " - pytree.ls[0] = Array(0, dtype=int32, weak_type=True)\n",
      " - pytree.ls[1] = Array(1, dtype=int32, weak_type=True)\n",
      " - pytree.ls[2] = Array(2, dtype=int32, weak_type=True)\n",
      " - pytree.x = 1.0\n",
      " - pytree.y = 42\n"
     ]
    }
   ],
   "source": [
    "class Bar(nnx.Pytree):\n",
    "  def __init__(self, x, use_bias: bool):\n",
    "    self.x = nnx.data(x)  # constrain inputs (e.g. user could pass Array or float)\n",
    "    self.y = nnx.data(42)  # force types that are not data by default\n",
    "    self.ls = nnx.List([jnp.array(i) for i in range(3)]) # use nnx.List for lists of data\n",
    "    self.bias = nnx.data(None)  # optional values that can be data\n",
    "    if use_bias:\n",
    "      self.bias = nnx.Param(jnp.array(0.0))\n",
    "\n",
    "pytree = Bar(1.0, True)\n",
    "pytree_structure(pytree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8055b72c",
   "metadata": {},
   "source": [
    "### Dataclasses\n",
    "`nnx.Pytree` dataclasses can be created by using the `nnx.dataclass` decorator. To control the status of each field, `nnx.static` and `nnx.data` can be used as `field` specifiers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9ca77d65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pytree structure:\n",
      " - pytree.ls[0].i = 0\n",
      " - pytree.ls[0].x = Array(0, dtype=int32, weak_type=True)\n",
      " - pytree.ls[1].i = 1\n",
      " - pytree.ls[1].x = Array(42, dtype=int32, weak_type=True)\n"
     ]
    }
   ],
   "source": [
    "import dataclasses\n",
    "\n",
    "@nnx.dataclass\n",
    "class Foo(nnx.Pytree):\n",
    "  i: int = nnx.data()\n",
    "  x: jax.Array\n",
    "  a: int\n",
    "  s: str = nnx.static(default='hi', kw_only=True)\n",
    "\n",
    "@nnx.dataclass\n",
    "class Bar(nnx.Pytree):\n",
    "  ls: list[Foo] = nnx.data()\n",
    "  shapes: list[int]\n",
    "\n",
    "pytree = Bar(\n",
    "  ls=[Foo(i, jnp.array(42 * i), hash(i)) for i in range(2)],\n",
    "  shapes=[8, 16, 32]\n",
    ")\n",
    "pytree_structure(pytree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca51f65",
   "metadata": {},
   "source": [
    "`dataclasses.dataclass` can also be used directly, however type checkers will not handle `nnx.static` and `nnx.data` correctly. To solve this `dataclasses.field` can be used by setting `metadata` with the appropriate entry for `static`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ff54e732",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataclass pytree structure:\n",
      " - pytree.a = 10\n"
     ]
    }
   ],
   "source": [
    "@dataclasses.dataclass\n",
    "class Bar(nnx.Pytree):\n",
    "  a: int = dataclasses.field(metadata={'static': False}) # data\n",
    "  b: str = dataclasses.field(metadata={'static': True})  # static\n",
    "\n",
    "pytree = Bar(a=10, b=\"hello\")\n",
    "pytree_structure(pytree, title='dataclass pytree structure')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6036a0e",
   "metadata": {},
   "source": [
    "### Attribute Updates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e0a1f63",
   "metadata": {},
   "source": [
    "The status of an attribute is defined during its first assignment and will not change upon reassignment. However, it is possible to override the status by explicitly using `nnx.data` or `nnx.static` on reassignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "509a517e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "original:\n",
      " - pytree.a = Array(1., dtype=float32, weak_type=True)\n",
      " - pytree.c = 3.14\n",
      "updated:\n",
      " - pytree.a = '🤔'\n",
      " - pytree.b = 42\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.a = jnp.array(1.0)  # data\n",
    "    self.b = \"Hello, world!\"               # static\n",
    "    self.c = nnx.data(3.14)         # data\n",
    "\n",
    "pytree = Foo()\n",
    "pytree_structure(pytree, \"original\")\n",
    "\n",
    "pytree.a = \"🤔\"  # data status doesn't change\n",
    "pytree.b = nnx.data(42)     # explicit annotation overrides status to data\n",
    "pytree.c = nnx.static(0.5)  # explicit annotation overrides status to static\n",
    "pytree_structure(pytree, \"updated\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17fd41f5",
   "metadata": {},
   "source": [
    "### Attribute checks\n",
    "`Pytree` has a variety of checks to prevent a common class of errors in JAX. This includes checking for Arrays being assigned to new `static` attributes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "98e04ff9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError: Found Arrays in value of type '<class 'jaxlib._jax.ArrayImpl'>' annotated with nnx.static(...) when setting attribute 'name' of Pytree type '<class '__main__.Foo'>'.\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree):\n",
    "  def __init__(self, name):\n",
    "    self.name = nnx.static(name)\n",
    "\n",
    "try:\n",
    "  foo = Foo(name=jnp.array(123))\n",
    "except ValueError as e:\n",
    "  print(\"ValueError:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070a07f0",
   "metadata": {},
   "source": [
    "Checking for Arrays being assigned to known `static` attributes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c864d5b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError: Cannot assign data value of type '<class 'jaxlib._jax.ArrayImpl'>' to static attribute 'name' of Pytree type '<class '__main__.Foo'>'. To override the status explicitly wrap the value with nnx.data on assignment:\n",
      "\n",
      "  _.name = nnx.data(...)\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "  foo = Foo(name=\"mattjj\")\n",
    "  foo.name = jnp.array(123)\n",
    "except ValueError as e:\n",
    "  print(\"ValueError:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3b17f07",
   "metadata": {},
   "source": [
    "Checking for Arrays after `__init__` on `static` attributes that could've been inserted via mutation. This check can be manually trigger via `nnx.check_pytree` at any time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "628698dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError: Found unexpected Arrays on value of type <class 'list'> in static attribute 'ls' of Pytree type '<class '__main__.Foo'>'. This is an error starting from Flax version 0.12.0.\n",
      "Consider one of the following options:\n",
      "\n",
      "1. If the attribute is meant to be static, either remove the Array value or wrap it in a static container.\n",
      "2. Wrap the value with nnx.data on assignment:\n",
      "\n",
      "  _.ls = nnx.data(...)\n",
      "\n",
      "3. Annotate the class attribute with nnx.Data:\n",
      "\n",
      "  class Foo(Pytree):\n",
      "    ls: nnx.Data[list]\n",
      "\n",
      "4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.\n",
      "5. Disable pytree for this class:\n",
      "\n",
      "  class Foo(Pytree, pytree=False):\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.ls = []  # treated as static\n",
    "    for i in range(5):\n",
    "      self.ls.append(jnp.array(i))  # append arrays into static attribute\n",
    "\n",
    "try:\n",
    "  foo = Foo()  # error: Array found in static attribute after `__init__`\n",
    "except ValueError as e:\n",
    "  print(\"ValueError:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37ee2429",
   "metadata": {},
   "source": [
    "Checking for `nnx.data` or `nnx.static` annotations stored inside nested structures that are not `nnx.Pytree` instances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f9d69634",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError: Found unexpected tags {'static', 'data'} on attribute 'Foo.a'. Values from nnx.data(...) and\n",
      "nnx.static(...) should be assigned to nnx.Pytree attributes directly, they should not be inside other structures. Got value of type '<class 'list'>' on Pytree of type '<class '__main__.Foo'>'.\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.a = [nnx.data(1), nnx.static(2)]  # annotations in sub-pytree\n",
    "\n",
    "try:\n",
    "  foo = Foo()\n",
    "except ValueError as e:\n",
    "  print(\"ValueError:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "452915ac",
   "metadata": {},
   "source": [
    "### Trace-level awareness\n",
    "To prevent tracer leakage NNX will raise an error when trying to update the attribute of a `Pytree` or the value of a `Variable` on instances that are passed as captures to functions called by JAX transforms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "668db479",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Cannot mutate 'Foo' from different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.count = nnx.data(0)\n",
    "\n",
    "foo = Foo()\n",
    "\n",
    "@jax.vmap  # or jit, grad, shard_map, pmap, scan, etc.\n",
    "def increment(n):\n",
    "  # foo passed as capture\n",
    "  foo.count += 1  # error!\n",
    "\n",
    "try:\n",
    "  increment(jnp.arange(5))\n",
    "except Exception as e:\n",
    "  print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "536d77b0",
   "metadata": {},
   "source": [
    "### Reference Sharing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2449a3c5",
   "metadata": {},
   "source": [
    "As the name implies Pytrees should be trees. To check if a structure is a well-defined tree you can use the `nnx.find_duplicates` functions which will return a list of duplicates, where each duplicate is a list of path tuples. In the example below we see that `left` and `right` are shared references therefore `find_duplicates` returns a non-empty list with the paths:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "32c46ce8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nnx.find_duplicates(m) = [[('left',), ('right',)]]  # not a tree\n"
     ]
    }
   ],
   "source": [
    "class Shared(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.x = jnp.array(1.0)\n",
    "\n",
    "class Parent(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.left = Shared()\n",
    "    self.right = self.left  # reference sharing\n",
    "\n",
    "m = Parent()\n",
    "\n",
    "print(f\"{nnx.find_duplicates(m) = }  # not a tree\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49b2e5f0",
   "metadata": {},
   "source": [
    "The main issue is that sharing is not preserved across pytree operations including JAX transforms, and this results in unintended state duplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c33e4862",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before: m.left is m.right = True\n",
      "Inside: m.left is m.right = False\n",
      "After:  m.left is m.right = False\n"
     ]
    }
   ],
   "source": [
    "m = Parent()\n",
    "print(f\"Before: {m.left is m.right = }\")\n",
    "\n",
    "@jax.jit\n",
    "def f(m):\n",
    "  print(f\"Inside: {m.left is m.right = }\")\n",
    "  return m\n",
    "\n",
    "m = f(m)\n",
    "print(f\"After:  {m.left is m.right = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d60953f4",
   "metadata": {},
   "source": [
    "Reference sharing is rare in most Machine Learning applications, however if it is required you can either use the `nnx.{split, merge, state, update}` APIs to move the deduplicated state and graph definiton across the JAX transforms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dda51b67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before: m.left is m.right = True\n",
      "Inside: m.left is m.right = True\n",
      "After:  m.left is m.right = True\n",
      "state = State({\n",
      "  'left': {\n",
      "    'x': Array(1., dtype=float32, weak_type=True)\n",
      "  }\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "m = Parent()\n",
    "print(f\"Before: {m.left is m.right = }\")\n",
    "graphdef, state = nnx.split(m)\n",
    "\n",
    "@jax.jit\n",
    "def f(graphdef, state):\n",
    "  m = nnx.merge(graphdef, state)\n",
    "  print(f\"Inside: {m.left is m.right = }\")\n",
    "  return nnx.state(m)\n",
    "\n",
    "state = f(graphdef, state)\n",
    "nnx.update(m, state)\n",
    "\n",
    "print(f\"After:  {m.left is m.right = }\")\n",
    "print(f\"{state = }\") # deduplicated state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0afee781",
   "metadata": {},
   "source": [
    "Or alternatively you can use the NNX transforms which preserve shared references:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "caa01e3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before: m.left is m.right = True\n",
      "Inside: m.left is m.right = True\n",
      "After:  m.left is m.right = True\n"
     ]
    }
   ],
   "source": [
    "m = Parent()\n",
    "print(f\"Before: {m.left is m.right = }\")\n",
    "\n",
    "@nnx.jit\n",
    "def f(m):\n",
    "  print(f\"Inside: {m.left is m.right = }\")\n",
    "  return m\n",
    "\n",
    "m = f(m)\n",
    "\n",
    "print(f\"After:  {m.left is m.right = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734aa5b3",
   "metadata": {},
   "source": [
    "### Turning off pytree registration\n",
    "`nnx.Pytree` allows you to turn off the pytree registration along with the attribute checks for subtypes by setting `pytree` type attribute option to `False`. This can be useful when upgrading to previous NNX code to newer Flax verions as you will still be able to use the NNX APIs or when creating types that should not be treated as pytree because e.g. they shared references."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d2e03753",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " nnx.state(foo) = State({\n",
      "  'a': {\n",
      "    0: Array(2, dtype=int32, weak_type=True),\n",
      "    1: Array(4, dtype=int32, weak_type=True)\n",
      "  },\n",
      "  'b': Array(6, dtype=int32, weak_type=True)\n",
      "})\n",
      " jax.tree_util.all_leaves([foo]) = True\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Pytree, pytree=False):\n",
    "  def __init__(self):\n",
    "    self.a = [jnp.array(1), jnp.array(2)]  # no checks\n",
    "    self.b = \"hello\" \n",
    "    self.b = jnp.array(3) # no checks\n",
    "\n",
    "foo = Foo()\n",
    "\n",
    "@nnx.jit # can use in NNX transformations\n",
    "def double(foo: Foo):\n",
    "  foo.a = [x * 2 for x in foo.a]\n",
    "  foo.b *= 2\n",
    "\n",
    "double(foo)\n",
    "print(f\"{ nnx.state(foo) = }\")  # can be used with NNX APIs\n",
    "print(f\"{ jax.tree_util.all_leaves([foo]) = }\")  # not a pytree"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfaf17de",
   "metadata": {},
   "source": [
    "## Module"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67000b88",
   "metadata": {},
   "source": [
    "NNX Modules are `Pytree`s that have two additional methods for traking intermediate values: `sow` and `perturb`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc5afc70",
   "metadata": {},
   "source": [
    "### sow\n",
    "`sow` receives a `Variable` type, a `name`, and a `value`, and stores it in the `Module` so it can be retrieved at a later time. As the following example shows, NNX APIs such as `nnx.state` or `nnx.pop` are a good way of retrieving the sowed state, however `pop` is recommended because it explicitly removes the temporary state from the Module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ca9f58a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "  \u001b[38;2;156;220;254m'blocks'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "    \u001b[38;2;156;220;254m0\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "      \u001b[38;2;156;220;254m'y_mean'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mIntermediate\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "        \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0mArray(4.659754e-06, dtype=float32),\u001b[38;2;255;213;3m)\u001b[0m\n",
      "      \u001b[38;2;255;213;3m)\u001b[0m\n",
      "    \u001b[38;2;255;213;3m}\u001b[0m,\n",
      "    \u001b[38;2;156;220;254m1\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "      \u001b[38;2;156;220;254m'y_mean'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mIntermediate\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "        \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0mArray(0.00025933, dtype=float32),\u001b[38;2;255;213;3m)\u001b[0m\n",
      "      \u001b[38;2;255;213;3m)\u001b[0m\n",
      "    \u001b[38;2;255;213;3m}\u001b[0m,\n",
      "    \u001b[38;2;156;220;254m2\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "      \u001b[38;2;156;220;254m'y_mean'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mIntermediate\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "        \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0mArray(0.05561922, dtype=float32),\u001b[38;2;255;213;3m)\u001b[0m\n",
      "      \u001b[38;2;255;213;3m)\u001b[0m\n",
      "    \u001b[38;2;255;213;3m}\u001b[0m\n",
      "  \u001b[38;2;255;213;3m}\u001b[0m\n",
      "\u001b[38;2;255;213;3m})\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "class Block(nnx.Module):\n",
    "  def __init__(self, din: int, dout: int, rngs: nnx.Rngs):\n",
    "    self.linear = nnx.Linear(din, dout, rngs=rngs)\n",
    "    self.bn = nnx.BatchNorm(dout, rngs=rngs)\n",
    "    self.dropout = nnx.Dropout(0.1, rngs=rngs)\n",
    "\n",
    "  def __call__(self, x):\n",
    "    y = nnx.relu(self.dropout(self.bn(self.linear(x))))\n",
    "    self.sow(nnx.Intermediate, \"y_mean\", jnp.mean(y))\n",
    "    return y\n",
    "\n",
    "class MLP(nnx.Module):\n",
    "  def __init__(self, num_layers, dim, rngs: nnx.Rngs):\n",
    "    self.blocks = nnx.List([Block(dim, dim, rngs) for _ in range(num_layers)])\n",
    "\n",
    "  def __call__(self, x):\n",
    "    for block in self.blocks:\n",
    "      x = block(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "model = MLP(num_layers=3, dim=20, rngs=nnx.Rngs(0))\n",
    "x = jnp.ones((10, 20))\n",
    "y = model(x)\n",
    "intermediates = nnx.pop(model, nnx.Intermediate) # extract intermediate values\n",
    "print(intermediates)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e266e5f",
   "metadata": {},
   "source": [
    "### perturb\n",
    "`perturb` is similar to `sow` but it aims to capture the gradient of a value, currently this is a two step process although it might be simplified in the future:\n",
    "1. Initialize the pertubation state by running the model once.\n",
    "2. Pass the perturbation state as a differentiable target to `grad`.\n",
    "\n",
    "As an example lets create a simple model and use `perturb` to get the intermediate gradient `xgrad` for the variable `x`, and initialize the perturbations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "41398e14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nnx.state(model, nnx.Perturbation) = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "  \u001b[38;2;156;220;254m'xgrad'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mPerturbation\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[0., 0., 0.]], dtype=float32)\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m})\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import optax\n",
    "\n",
    "class Model(nnx.Module):\n",
    "  def __init__(self, rngs):\n",
    "    self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n",
    "    self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n",
    "  def __call__(self, x):\n",
    "    x = nnx.gelu(self.linear1(x))\n",
    "    x = self.perturb('xgrad', x)\n",
    "    x = self.linear2(x)\n",
    "    return x\n",
    "\n",
    "rngs = nnx.Rngs(0)\n",
    "model = Model(rngs)\n",
    "optimizer = nnx.Optimizer(model, tx=optax.sgd(1e-1), wrt=nnx.Param)\n",
    "x, y = rngs.uniform((1, 2)), rngs.uniform((1, 4))\n",
    "_ = model(x) # initialize perturbations\n",
    "print(f\"{nnx.state(model, nnx.Perturbation) = !s}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9221005",
   "metadata": {},
   "source": [
    "Next we'll create a training step function that differentiates w.r.t. both the parameters of the model and the perturbations, the later will be the gradients for the intermediate values. `nnx.jit` and `nnx.value_and_grad` will be use to automatically propagate state updates. We'll return the `loss` function and the itermediate gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d10effba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step = 0, loss = Array(0.7326511, dtype=float32), iterm_grads = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "  \u001b[38;2;156;220;254m'xgrad'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mPerturbation\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[-0.430146  , -0.14356601,  0.2935633 ]], dtype=float32)\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m})\u001b[0m\n",
      "step = 1, loss = Array(0.65039134, dtype=float32), iterm_grads = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "  \u001b[38;2;156;220;254m'xgrad'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mPerturbation\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[-0.38535568, -0.11745065,  0.24441527]], dtype=float32)\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m})\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "@nnx.jit\n",
    "def train_step(model, optimizer, x, y):\n",
    "  graphdef, params, perturbations = nnx.split(model, nnx.Param, nnx.Perturbation)\n",
    "\n",
    "  def loss_fn(params, perturbations):\n",
    "    model = nnx.merge(graphdef, params, perturbations)\n",
    "    return jnp.mean((model(x) - y) ** 2)\n",
    "\n",
    "  loss, (grads, iterm_grads) = nnx.value_and_grad(loss_fn, argnums=(0, 1))(params, perturbations)\n",
    "  optimizer.update(model, grads)\n",
    "\n",
    "  return loss, iterm_grads\n",
    "\n",
    "for step in range(2):\n",
    "  loss, iterm_grads = train_step(model, optimizer, x, y)\n",
    "  print(f\"{step = }, {loss = }, {iterm_grads = !s}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8511c3d",
   "metadata": {},
   "source": [
    "## Object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2fd61c",
   "metadata": {},
   "source": [
    "`Object` are NNX types that are **not** registered as JAX pytrees. Formally, any `Object` subclass is a `nnx.Pytree` with `pytree=False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a9cab639",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " nnx.state(foo) = State({\n",
      "  'a': {\n",
      "    0: Array(2, dtype=int32, weak_type=True),\n",
      "    1: Array(4, dtype=int32, weak_type=True)\n",
      "  },\n",
      "  'b': Array(6, dtype=int32, weak_type=True)\n",
      "})\n",
      " jax.tree_util.all_leaves([foo]) = True\n"
     ]
    }
   ],
   "source": [
    "class Foo(nnx.Object): # instead of Foo(nnx.Pytree, pytree=False)\n",
    "  def __init__(self):\n",
    "    self.a = [jnp.array(1), jnp.array(2)]  # no checks\n",
    "    self.b = \"hello\" \n",
    "    self.b = jnp.array(3) # no checks\n",
    "\n",
    "foo = Foo()\n",
    "\n",
    "@nnx.jit # can use in NNX transformations\n",
    "def double(foo: Foo):\n",
    "  foo.a = [x * 2 for x in foo.a]\n",
    "  foo.b *= 2\n",
    "\n",
    "double(foo)\n",
    "print(f\"{ nnx.state(foo) = }\")  # can be used with NNX APIs\n",
    "print(f\"{ jax.tree_util.all_leaves([foo]) = }\")  # not a pytree"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
